Pin jax to 0.11.0 in the three execution-* workflows - #620
Conversation
The 2026-08-19 pin (#617) covered cache.yml, ci.yml and publish.yml, but this repo has six workflows that install jax. execution-linux.yml, execution-osx.yml and execution-win.yml install "jax[CPU]" unpinned and build with `jb build lectures -W --keep-going`, which executes every lecture -- so jax 0.11.1 fails them on CellTimeoutError in numpy_vs_numba_vs_jax. execution-linux.yml runs on a daily cron and has been red since 2026-08-18: run 32152771492 (08-18) and run 32268416676 (08-19, head_sha 915bfd9, i.e. after #617 merged) both log "Successfully installed jax-0.11.1 jaxlib-0.11.1" and then CellTimeoutError at the 600s myst-nb timeout. The 08-16 run resolved jax 0.11.0 and executed the same lecture in 4.35 seconds. Pinning jax alone is sufficient: jax 0.11.0's metadata carries Requires-Dist: jaxlib<=0.11.0,>=0.11.0, and the pairing is additionally enforced at import in both directions, so jaxlib cannot drift. Closes #619 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Pins JAX to 0.11.0 in the remaining three “execution-*” GitHub Actions workflows so the full jb build ... -W execution checks don’t pick up the jax/jaxlib 0.11.1 regression and fail on CellTimeoutError. This aligns these workflows with the existing pins already present in cache.yml, ci.yml, and publish.yml (issue #619).
Changes:
- Pin
pip install "jax[CPU]==0.11.0"inexecution-linux.yml,execution-osx.yml, andexecution-win.yml. - Add inline rationale comments in each workflow explaining the regression and why the pin is needed.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| .github/workflows/execution-linux.yml | Pin CPU JAX to 0.11.0 for the daily execution build to avoid the 0.11.1 regression. |
| .github/workflows/execution-osx.yml | Pin CPU JAX to 0.11.0 for the scheduled macOS execution build to avoid the 0.11.1 regression. |
| .github/workflows/execution-win.yml | Pin CPU JAX to 0.11.0 for the scheduled Windows execution build to avoid the 0.11.1 regression. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
Verified by dispatching Execution Tests [Latest Anaconda, Linux] on this branch: run 32320365735, head_sha The green tick is not the evidence — this one is, because it shows the lecture was actually executed rather than restored from cache:
Whole-run figures: 52 notebooks executed, 357.0 s aggregate, zero execution failures, zero For comparison against the same workflow before this change: run 32268416676 (2026-08-19) logged One trap for whoever audits this next: a plain |
The wrap split `!pip install quantecon jax` across two comment lines, so skimming the first showed `!pip install quantecon` with `jax` orphaned behind the next `# ` -- dropping the one token the sentence exists to explain. Matches the unbroken form already used in cache.yml, ci.yml and publish.yml. Comments only; the pin is unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The workflow pin from #620 protects the build but not the reader. In CI the cell is a no-op because jax 0.11.0 is already installed, so pip treats the unpinned requirement as satisfied. In Colab nothing pre-satisfies it: a reader executing this notebook today resolves jax 0.11.1, then reaches the CPU-pinned lax.fori_loop and lax.scan cells at n = 10,000,000. jax 0.11.1 regresses XLA:CPU execution quadratically -- measured on linux x86_64, doubling n multiplies runtime by ~4.1 (fori) and ~4.0 (scan). At n=400,000 it is 105s against 0.005s under 0.11.0; extrapolated to the lecture's n=10,000,000 that is roughly 18 hours, which a reader experiences as a hang. Pinned to ==0.11.0 rather than !=0.11.1 deliberately: the regression is still present on jax main (nightly 0.11.2.dev20260819 measured at 96% of 0.11.1's time), so an exclusion would admit a likely-broken 0.11.2 without warning. This matches the exact version the six workflows install, so readers now run what CI runs. The other two jax cells in this repo are untouched: jax_intro.md loops to n=20 and autodiff.md has no lax loop, so neither reaches the regime. Lifting this pin is tracked alongside the workflow pins in QuantEcon/workspace-lectures#49. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Pins jax to
0.11.0inexecution-linux.yml,execution-osx.ymlandexecution-win.yml, matching the pin already carried bycache.yml,ci.ymlandpublish.yml.Closes #619
Why
The 2026-08-19 pin in #617 covered three of the six workflows in this repo that install jax. The other three install
jax[CPU]with no constraint and build withjb build lectures --path-output=./ -W --keep-going, which executes every lecture — so they hit the jax 0.11.1 regression head-on.execution-linux.ymlruns on a daily cron and has been failing since 2026-08-18:Successfully installed jax-0.11.0 jaxlib-0.11.0→numpy_vs_numba_vs_jax.md: Executed notebook in 4.35 seconds→ successsetup-miniconda, jax never installedSuccessfully installed jax-0.11.1 jaxlib-0.11.1→CellTimeoutErrorhead_sha 915bfd9f— i.e. after #617 mergedexecution-osx.yml(Mondays) andexecution-win.yml(Thursdays) run the same unpinned install and the same-Wbuild, so they are exposed on their own schedules.Those two 0.11.0-vs-0.11.1 runs are incidentally the cleanest controlled experiment available for this regression: same workflow, same runner image, only the jax version differing — 4.35 s versus a 600 s timeout.
Notes on the pin
Pinning jax alone is sufficient.
jax==0.11.0's wheel metadata carriesRequires-Dist: jaxlib<=0.11.0,>=0.11.0, and the pairing is additionally enforced at import in both directions (jaxlib version 0.11.1 is newer than and incompatible with jax version 0.11.0), so jaxlib cannot drift even under--no-deps.jax[CPU]is kept as-is rather than simplified tojax:cpuis a declared extra on jax 0.11.0 that adds no requirements, so the bracket is a harmless no-op and dropping it would be an unrelated change.On the comment wording
The new comments describe the regression as XLA:CPU execution going quadratic, rather than as a
lax.fori_loophang. That is a correction established while validating the settle week (QuantEcon/workspace-lectures#51): re-run on linux x86_64, native linux aarch64 and macOS arm64, the stall is a ~O(n²) blow-up at runtime — a stack sample sits inxla::cpu::ThunkExecutor::ExecuteSequential, so compilation has already finished — andlax.scanregresses identically tolax.fori_loop. Measured on linux x86_64 at n=100k/200k/400k:fori6.27 s / 28.49 s / 103.26 s andscan3.83 s / 15.45 s / 60.43 s under 0.11.1, against a flat ~0.05 s under 0.11.0.The three previously-pinned workflows still carry the older
device=cpuwording, which has been stale since #563 removed that argument on 2026-06-19. Left alone here to keep this diff to the pin; worth a separate tidy-up.Verifying
workflow_dispatchis enabled on all three, so a manual dispatch of Execution Tests [Latest Anaconda, Linux] on this branch is the direct check — expectSuccessfully installed jax-0.11.0andnumpy_vs_numba_vs_jax.md: Executed notebook in N secondsat single-digit seconds, with noCellTimeoutError.The unpin condition for all six workflows is tracked on QuantEcon/workspace-lectures#49.